import os
# 设置环境变量，允许重复加载OpenMP库
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
from ultralytics import YOLO

model = YOLO("./yolov8n.pt")  # 加载预训练模型（建议用于训练）

if __name__ == '__main__':
    #imgsz=640输入大小
    model.train(data="./dataset/mydata.yaml", epochs=3, imgsz=640)  # 训练模型

    # metrics = model.val()  # 在验证集上评估模型性能
    # success = model.export(format="onnx")  # 将模型导出为 ONNX 格式
    # results = model("./dataset/test/images/img0.jpg")  # 对图像进行预测
    # # Process results list
    # for result in results:
    #     boxes = result.boxes  # Boxes object for bounding box outputs
    #     masks = result.masks  # Masks object for segmentation masks outputs
    #     keypoints = result.keypoints  # Keypoints object for pose outputs
    #     probs = result.probs  # Probs object for classification outputs
    #     result.show()  # display to screen